{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangGraph Sidekick – Week 4 Final Assignment\n",
    "\n",
    "## Author: Bharat Puri\n",
    "\n",
    "## Objective:\n",
    "\n",
    "Extend and personalize your LangGraph-based “Sidekick” by adding custom tools, persistent SQL memory, and a planning mechanism.\n",
    "The assistant should demonstrate autonomy, context retention, and the ability to use tools intelligently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U -q imapclient langchain==1.0.2 langchain-openai==1.0.1 langchain-chroma==1.0.0 langchain-community==0.4 langchain-core==1.0.0 langchain-text-splitters==1.0.0 langchain-huggingface==1.0.0 langchain-classic==1.0.0 chromadb==1.2.1 sentence-transformers==5.1.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q langgraph sqlite-utils gradio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 1. Imports\n",
    "# -------------------------------------------------------------\n",
    "from dotenv import load_dotenv\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_community.chat_message_histories import SQLChatMessageHistory\n",
    "from langgraph.graph import StateGraph, END\n",
    "import gradio as gr\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Environment loaded successfully.\n"
     ]
    }
   ],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 2. Load Environment Variables\n",
    "# -------------------------------------------------------------\n",
    "load_dotenv(override=True)\n",
    "\n",
    "# Required in .env:\n",
    "# OPENAI_API_KEY=sk-xxxx\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    raise EnvironmentError(\"❌ OPENAI_API_KEY not found. Please check your .env file.\")\n",
    "else:\n",
    "    print(\"✅ Environment loaded successfully.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 3. Define Custom Tools\n",
    "# -------------------------------------------------------------\n",
    "@tool\n",
    "def search_docs(query: str, folder: str = \"./docs\"):\n",
    "    \"\"\"\n",
    "    Search .txt files in a folder for a given keyword.\n",
    "    \"\"\"\n",
    "    results = []\n",
    "    if not os.path.exists(folder):\n",
    "        return f\"❌ Folder '{folder}' not found.\"\n",
    "\n",
    "    for file in os.listdir(folder):\n",
    "        if file.endswith(\".txt\"):\n",
    "            with open(os.path.join(folder, file), \"r\", encoding=\"utf-8\") as f:\n",
    "                text = f.read().lower()\n",
    "                if query.lower() in text:\n",
    "                    results.append(file)\n",
    "    return results or \"No matches found.\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def run_python(expression: str):\n",
    "    \"\"\"\n",
    "    Safely evaluate simple Python expressions.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        result = eval(expression, {\"__builtins__\": {}})\n",
    "        return f\"✅ Result: {result}\"\n",
    "    except Exception as e:\n",
    "        return f\"❌ Error: {e}\"\n",
    "\n",
    "tools = [search_docs, run_python]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 4. Persistent SQL Memory Setup\n",
    "# -------------------------------------------------------------\n",
    "memory = SQLChatMessageHistory(\n",
    "    connection_string=\"sqlite:///sidekick_memory.db\",\n",
    "    session_id=\"bharat_user\"\n",
    ")\n",
    "\n",
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "\n",
    "def remember_message(role: str, message: str):\n",
    "    \"\"\"\n",
    "    Stores user/assistant messages in persistent SQL memory (LangChain v1.0+ format).\n",
    "    \"\"\"\n",
    "    if role == \"user\":\n",
    "        memory.add_message(HumanMessage(content=message))\n",
    "    elif role == \"assistant\":\n",
    "        memory.add_message(AIMessage(content=message))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 5. Initialize LangGraph Agents\n",
    "# -------------------------------------------------------------\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "\n",
    "planner_prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "You are the PLANNER agent. Analyze the user request and describe what tasks should be performed.\n",
    "User request: {user_input}\n",
    "\"\"\")\n",
    "\n",
    "worker_prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "You are the WORKER agent. Follow the planner's instruction carefully.\n",
    "Instruction: {instruction}\n",
    "\"\"\")\n",
    "\n",
    "def planner_node(state):\n",
    "    plan = llm.invoke(planner_prompt.format(user_input=state[\"user_input\"]))\n",
    "    return {\"instruction\": plan.content}\n",
    "\n",
    "def worker_node(state):\n",
    "    result = llm.invoke(worker_prompt.format(instruction=state[\"instruction\"]))\n",
    "    return {\"result\": result.content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 6. Build LangGraph Workflow (with state_schema)\n",
    "# -------------------------------------------------------------\n",
    "from typing import TypedDict\n",
    "\n",
    "# Define the structure of your graph state\n",
    "class GraphState(TypedDict):\n",
    "    user_input: str\n",
    "    instruction: str\n",
    "    result: str\n",
    "\n",
    "# Initialize the state graph with schema\n",
    "graph = StateGraph(GraphState)\n",
    "\n",
    "# Add planner and worker nodes\n",
    "graph.add_node(\"planner\", planner_node)\n",
    "graph.add_node(\"worker\", worker_node)\n",
    "\n",
    "# Connect the flow planner → worker → END\n",
    "graph.add_edge(\"planner\", \"worker\")\n",
    "graph.add_edge(\"worker\", END)\n",
    "\n",
    "# Define the entry point\n",
    "graph.set_entry_point(\"planner\")\n",
    "\n",
    "# Compile the graph\n",
    "compiled_graph = graph.compile()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 7. Clarifying Questions\n",
    "# -------------------------------------------------------------\n",
    "def ask_clarifying_questions(user_input: str) -> str:\n",
    "    \"\"\"\n",
    "    Generates clarifying questions to ensure task understanding.\n",
    "    \"\"\"\n",
    "    clarifications = [\n",
    "        f\"Can you clarify what exactly you want me to do with: '{user_input}'?\",\n",
    "        \"Should I use a specific folder or tool for this task?\",\n",
    "        \"What level of detail do you expect in the output?\"\n",
    "    ]\n",
    "    return \"\\n\".join(clarifications)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 8. Core Assistant Logic\n",
    "# -------------------------------------------------------------\n",
    "def run_sidekick(user_input: str):\n",
    "    \"\"\"\n",
    "    Executes the LangGraph assistant flow.\n",
    "    \"\"\"\n",
    "    print(\"🤖 Step 1: Asking clarifying questions...\")\n",
    "    clarifications = ask_clarifying_questions(user_input)\n",
    "\n",
    "    print(\"🧠 Step 2: Planning and executing tasks...\")\n",
    "    remember_message(\"user\", user_input)\n",
    "    result = compiled_graph.invoke({\"user_input\": user_input})\n",
    "    remember_message(\"assistant\", str(result))\n",
    "\n",
    "    return f\"### Clarifications\\n{clarifications}\\n\\n### Response\\n{result['result']}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 9. Gradio Interface\n",
    "# -------------------------------------------------------------\n",
    "def interface_fn(user_input):\n",
    "    return run_sidekick(user_input)\n",
    "\n",
    "iface = gr.Interface(\n",
    "    fn=interface_fn,\n",
    "    inputs=gr.Textbox(\n",
    "        label=\"Enter your request\",\n",
    "        placeholder=\"e.g., Search for AI papers or evaluate 3*8+5\",\n",
    "        lines=2,\n",
    "    ),\n",
    "    outputs=gr.Markdown(label=\"AI Sidekick Response\"),\n",
    "    title=\"🧠 LangGraph AI Sidekick with SQL Memory\",\n",
    "    description=(\n",
    "        \"A personalized assistant powered by LangGraph and OpenAI.\\n\\n\"\n",
    "        \"💡 Supports planning, task execution, and SQL-based persistent memory.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 9b. Enhanced Gradio Interface with Memory Viewer\n",
    "# -------------------------------------------------------------\n",
    "import pandas as pd\n",
    "\n",
    "def view_memory():\n",
    "    \"\"\"\n",
    "    Parses and displays recent messages from SQL memory in a chat-style format.\n",
    "    Works with latest LangChain SQLChatMessageHistory schema.\n",
    "    \"\"\"\n",
    "    import sqlite3, json\n",
    "\n",
    "    try:\n",
    "        conn = sqlite3.connect(\"sidekick_memory.db\")\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute(\"SELECT message FROM message_store ORDER BY id DESC LIMIT 10\")\n",
    "        rows = cursor.fetchall()\n",
    "        conn.close()\n",
    "\n",
    "        if not rows:\n",
    "            return \"🗂️ No messages found in memory.\"\n",
    "\n",
    "        messages_md = \"### 🧠 Conversation Memory\\n\\n\"\n",
    "\n",
    "        for row in reversed(rows):\n",
    "            try:\n",
    "                msg_obj = json.loads(row[0])\n",
    "                msg_type = msg_obj.get(\"type\", \"\")\n",
    "                msg_data = msg_obj.get(\"data\", {})\n",
    "                content = msg_data.get(\"content\", \"\")\n",
    "                if msg_type == \"human\":\n",
    "                    messages_md += f\"🧑 **User:** {content}\\n\\n\"\n",
    "                elif msg_type == \"ai\":\n",
    "                    messages_md += f\"🤖 **Assistant:** {content}\\n\\n\"\n",
    "                else:\n",
    "                    messages_md += f\"💬 **{msg_type.capitalize()}:** {content}\\n\\n\"\n",
    "            except Exception as inner_e:\n",
    "                messages_md += f\"⚠️ [Unreadable Entry] {row[0][:80]}...\\n\\n\"\n",
    "\n",
    "        return messages_md\n",
    "\n",
    "    except Exception as e:\n",
    "        return f\"⚠️ Error reading memory: {e}\"\n",
    "\n",
    "\n",
    "# Main tab – Chat interface\n",
    "with gr.Blocks() as iface:\n",
    "    gr.Markdown(\"## 🧠 LangGraph AI Sidekick with SQL Memory\")\n",
    "    gr.Markdown(\"A personalized assistant powered by LangGraph and OpenAI.\\n\\n\"\n",
    "                \"💡 Supports planning, task execution, and SQL-based persistent memory.\")\n",
    "    \n",
    "    with gr.Tab(\"💬 Chat\"):\n",
    "        input_box = gr.Textbox(\n",
    "            label=\"Enter your request\",\n",
    "            placeholder=\"e.g., Search for AI papers or evaluate 3*8+5\",\n",
    "            lines=2,\n",
    "        )\n",
    "        output_box = gr.Markdown(label=\"AI Sidekick Response\")\n",
    "        submit_btn = gr.Button(\"Submit\", variant=\"primary\")\n",
    "        submit_btn.click(fn=interface_fn, inputs=input_box, outputs=output_box)\n",
    "\n",
    "    with gr.Tab(\"🗂️ Memory Viewer\"):\n",
    "        gr.Markdown(\"### Recent Conversation Memory\")\n",
    "        memory_output = gr.Markdown()\n",
    "        refresh_btn = gr.Button(\"🔄 Refresh Memory\")\n",
    "        refresh_btn.click(fn=view_memory, inputs=None, outputs=memory_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7863\n",
      "* To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7863/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖 Step 1: Asking clarifying questions...\n",
      "🧠 Step 2: Planning and executing tasks...\n"
     ]
    }
   ],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 10. Launch App\n",
    "# -------------------------------------------------------------\n",
    "iface.launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
